#!/usr/bin/env python

import sys

try:
    import OpenBayes
except Exception, e:
    print('This example requires OpenBayes. Please Visit http://www.openbayes.org/ for install instructions')
    sys.exit(-1)

class NaiveBayesianNetworkOpenBayes(object):
    """
        Creates a Bayesian Network using the DAG supplied by OpenBayes
    """
    def __init__(self):
        self.graph = OpenBayes.BNet('Bayesian Classification Toy Model')

    def build_root(self,rootName):
        self.rootName = rootName
        self.root = self.graph.add_v(OpenBayes.BVertex(rootName, True, 2))
        
    def build_topology(self,attributes):
        #topology buiding
        
        self.attributes = {}
        self.attributesName = attributes
        
        for attName in attributes:
            att = self.graph.add_v( OpenBayes.BVertex(attName, True, 2))
            self.graph.add_e( OpenBayes.DirEdge(len(self.graph.e), self.root, att))
            self.attributes[attName] = att
        
        self.graph.InitDistributions()
            
    def set_root_apriori_distribution(self,values):
        self.root.setDistributionParameters(values)

    def train(self,trainingSet):
        
        parameters = {}
        
        totalSpamDocs = float(len(trainingSet['spam']))
        totalNoSpamDocs = float(len(trainingSet['noSpam']))
        
        spamFrecuencies = {}
        noSpamFrecuencies = {}

        for key in self.attributes.keys():
            spamFrecuencies[key] = float(0)
            noSpamFrecuencies[key] = float(0)
        
        wordsInSpam = 0 
        wordsInNoSpam = 0
        
        for document in trainingSet['spam']:
            for key in document.keys():
                spamFrecuencies[key] = spamFrecuencies[key] + float(document[key])
                wordsInSpam = wordsInSpam + float(document[key])
        
        for document in trainingSet['noSpam']:
            for key in document.keys():
                noSpamFrecuencies[key] = noSpamFrecuencies[key] + float(document[key])
                wordsInNoSpam = wordsInNoSpam + float(document[key])              

        print('P(C=0) = '+ str(float(wordsInNoSpam)/float(wordsInNoSpam+wordsInSpam)))
        print('P(C=1) = '+ str(float(wordsInSpam)/float(wordsInNoSpam+wordsInSpam)))
        
        self.set_root_apriori_distribution([float(wordsInNoSpam)/float(wordsInNoSpam+wordsInSpam),float(wordsInSpam)/float(wordsInNoSpam+wordsInSpam)])
        
        for attName in self.attributes.keys():

            self.attributes[attName].setDistributionParameters([noSpamFrecuencies[attName]/wordsInNoSpam,1-spamFrecuencies[attName]/wordsInSpam,1-noSpamFrecuencies[attName]/wordsInNoSpam,spamFrecuencies[attName]/wordsInSpam])
            print('----------------------------------------')
            print('Si no es spam')
            print(noSpamFrecuencies[attName]/wordsInNoSpam)
            print(1-spamFrecuencies[attName]/wordsInSpam)
            print('Si es spam')
            print(1-noSpamFrecuencies[attName]/wordsInNoSpam)
            print(spamFrecuencies[attName]/wordsInSpam)
            print('Cierre')
            print(noSpamFrecuencies[attName]/wordsInNoSpam + (1-noSpamFrecuencies[attName]/wordsInNoSpam))
            print(1-spamFrecuencies[attName]/wordsInSpam + (spamFrecuencies[attName]/wordsInSpam))
            print('----------------------------------------')
    
    def classify(self,doc):

        ie = OpenBayes.MCMCEngine(self.graph)
        ie.SetObs(doc)
        result = ie.Marginalise(self.rootName)
        return result
        
class NaiveBayesianNetworkCustom(object):
    """
        Creates a Bayesian Network using the DAG supplied by OpenBayes
    """
    def __init__(self):
        self.verbose = False
        
    def set_verbose_mode(self):
        self.verbose = True
        
    def build_root(self,rootName):
        self.rootName = rootName

    def build_topology(self,attributes):
        #topology buiding
        self.conditionalProbabilities = {}

        self.attributes = {}
        self.attributesName = attributes
        
        for attName in attributes:
            self.attributes[attName] = attName

    def set_root_apriori_distribution(self,values):
            
        self.aPrioriProbabilities = values 

    def train(self,trainingSet):

        parameters = {}

        totalSpamDocs = float(len(trainingSet['spam']))
        totalNoSpamDocs = float(len(trainingSet['noSpam']))

        spamFrecuencies = {}
        noSpamFrecuencies = {}

        for key in self.attributes.keys():
            spamFrecuencies[key] = float(0)
            noSpamFrecuencies[key] = float(0)

        wordsInSpam = 0 
        wordsInNoSpam = 0

        for document in trainingSet['spam']:
            for key in document.keys():
                if (document[key] == 1):
                    spamFrecuencies[key] = spamFrecuencies[key] + 1

        for document in trainingSet['noSpam']:
            for key in document.keys():
                if (document[key] == 0):
                    noSpamFrecuencies[key] = noSpamFrecuencies[key] + 1
        
        if (self.verbose):
            print('----------------------------------------')
            print(float(totalSpamDocs)/float(totalSpamDocs+totalNoSpamDocs))
            print(float(totalNoSpamDocs)/float(totalSpamDocs+totalNoSpamDocs))
            
        self.set_root_apriori_distribution([float(totalNoSpamDocs)/float(totalSpamDocs+totalNoSpamDocs),float(totalSpamDocs)/float(totalSpamDocs+totalNoSpamDocs)])

        for attName in self.attributes.keys():
            self.conditionalProbabilities[attName] = {}
            self.conditionalProbabilities[attName][0] = {} 
            self.conditionalProbabilities[attName][1] = {}
            #format [attribute_name][attribute=value][category=value]
            self.conditionalProbabilities[attName][0][0] = float(noSpamFrecuencies[attName]) / float(totalNoSpamDocs)
            self.conditionalProbabilities[attName][0][1] = 1 - float(spamFrecuencies[attName]) / float(totalSpamDocs)
            self.conditionalProbabilities[attName][1][0] = 1 - float(noSpamFrecuencies[attName]) / float(totalNoSpamDocs)
            self.conditionalProbabilities[attName][1][1] = float(spamFrecuencies[attName]) / float(totalSpamDocs)
            
            #sample correction to prevent 0 probability information wipeout
            
            if (self.conditionalProbabilities[attName][0][0] == 0.0):
                self.conditionalProbabilities[attName][0][0] = self.conditionalProbabilities[attName][0][0] + 0.0000001
                self.conditionalProbabilities[attName][1][0] = self.conditionalProbabilities[attName][1][0] - 0.0000001

            if (self.conditionalProbabilities[attName][1][0] == 0.0):
                self.conditionalProbabilities[attName][1][0] = self.conditionalProbabilities[attName][1][0] + 0.0000001
                self.conditionalProbabilities[attName][0][0] = self.conditionalProbabilities[attName][0][0] - 0.0000001

            if (self.conditionalProbabilities[attName][0][1] == 0.0):
                self.conditionalProbabilities[attName][0][1] = self.conditionalProbabilities[attName][0][1] + 0.0000001
                self.conditionalProbabilities[attName][1][1] = self.conditionalProbabilities[attName][1][1] - 0.0000001

            if (self.conditionalProbabilities[attName][1][1] == 0.0):
                self.conditionalProbabilities[attName][1][1] = self.conditionalProbabilities[attName][1][1] + 0.0000001
                self.conditionalProbabilities[attName][0][1] = self.conditionalProbabilities[attName][0][1] - 0.0000001


            if (self.verbose):
                print('----------------------------------------')
                print('P(' + attName + '= 0 | C = 0) = ' + str(self.conditionalProbabilities[attName][0][0]))
                print('P(' + attName + '= 0 | C = 1) = ' + str(self.conditionalProbabilities[attName][0][1]))
                print('P(' + attName + '= 1 | C = 0) = ' + str(self.conditionalProbabilities[attName][1][0]))
                print('P(' + attName + '= 1 | C = 1) = ' + str(self.conditionalProbabilities[attName][1][1]))
                print('Cierre')
                print(self.conditionalProbabilities[attName][0][0] + self.conditionalProbabilities[attName][1][0])
                print(self.conditionalProbabilities[attName][0][1] + self.conditionalProbabilities[attName][1][1])
                print('----------------------------------------')


    def classify(self,doc):

        noSpamProbability = self.aPrioriProbabilities[0]
        spamProbability = self.aPrioriProbabilities[1]

        for att in self.attributesName:
            noSpamProbability = noSpamProbability * self.conditionalProbabilities[att][doc[att]][0]

        for att in self.attributesName:
            spamProbability = spamProbability * self.conditionalProbabilities[att][doc[att]][1]

        return [noSpamProbability, spamProbability]         
